from neo4j import GraphDatabase
import csv
import json

def sanitize_label(label):
    """清理标签，移除非法字符"""
    return label.strip().replace("/", "_").replace(" ", "_").replace("'", "").replace('"', '').replace("@", "AT_").replace(":", "_")

def import_nodes(uri, user, password, nodes_input_file, batch_size=1000):
    driver = GraphDatabase.driver(uri, auth=(user, password))
    with driver.session() as session:
        with open(nodes_input_file, "r", newline='', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            batch = []
            for row_num, row in enumerate(reader, 1):
                try:
                    # 处理id字段：去除首尾空格并转换为整数
                    raw_id = row["id"].strip()
                    if not raw_id:
                        raise ValueError("id字段为空")
                    node_id = int(raw_id)

                    labels = [sanitize_label(label) for label in row["labels"].strip("[]").split(", ")]
                    properties = json.loads(row["properties"]) if row["properties"] else {}

                    # 将id改为nid
                    properties["nid"] = node_id

                    batch.append((labels, properties))

                    if len(batch) >= batch_size:
                        save_to_csv(batch, "file/batchNodes.csv")
                        batch = []

                except Exception as e:
                    print(f"处理行 {row_num} 时出错: {str(e)}")
                    print(f"原始数据: {row}")
                    with open("error_log.txt", "a", encoding='utf-8') as log:
                        log.write(f"处理行 {row_num} 时出错: {str(e)}\n")
                        log.write(f"数据: {row}\n")
                    continue  # 跳过当前行

            if batch:
                save_to_csv(batch, "file/batchNodes.csv")

    print("节点导入完成。")

def import_relationships(uri, user, password, relationships_input_file, batch_size=1000):
    driver = GraphDatabase.driver(uri, auth=(user, password))
    with driver.session() as session:
        with open(relationships_input_file, "r", newline='', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            batch = []
            for row_num, row in enumerate(reader, 1):
                try:
                    # 处理start_id和end_id字段：去除首尾空格并转换为整数
                    start_id = int(row["start_id"].strip())
                    end_id = int(row["end_id"].strip())

                    rel_type = sanitize_label(row["type"].strip())
                    properties = json.loads(row["properties"]) if row["properties"] else {}

                    batch.append((start_id, end_id, rel_type, properties))

                    if len(batch) >= batch_size:
                        save_to_csv(batch, "file/batchRelationships.csv")
                        batch = []

                except Exception as e:
                    print(f"处理行 {row_num} 时出错: {str(e)}")
                    print(f"原始数据: {row}")
                    with open("error_log.txt", "a", encoding='utf-8') as log:
                        log.write(f"处理行 {row_num} 时出错: {str(e)}\n")
                        log.write(f"数据: {row}\n")
                    continue  # 跳过当前行

            if batch:
                save_to_csv(batch, "file/batchRelationships.csv")

    print("关系导入完成。")

def save_to_csv(batch, file_name):
    with open(file_name, "a", newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        if "Nodes" in file_name:
            for labels, properties in batch:
                labels_str = ", ".join(labels)
                properties_str = json.dumps(properties)
                writer.writerow([labels_str, properties_str])
        elif "Relationships" in file_name:
            for start_id, end_id, rel_type, properties in batch:
                properties_str = json.dumps(properties)
                writer.writerow([start_id, end_id, rel_type, properties_str])

if __name__ == "__main__":
    neo4j_uri = "bolt://localhost:7687"
    neo4j_user = "neo4j"
    neo4j_password = "88888888"
    nodes_input_file = "import/nodes.csv"
    relationships_input_file = "import/relationships.csv"

    import_nodes(neo4j_uri, neo4j_user, neo4j_password, nodes_input_file)
    import_relationships(neo4j_uri, neo4j_user, neo4j_password, relationships_input_file)
